import re
import os
import requests
import pandas as pd
import time
import json
from typing import Dict, List, Optional
import pymysql
from pymysql import Error

class CPICrawler:
    """国家统计局CPI数据采集器"""
    
    # 类常量：指标映射字典
    # 对应的key是在国家数据网站上 通过开发者工具 分析每个阶段的xhr中的负载部分确定的
    CODE_MAPPING = {
        'A01030G': {
            'A01030G01': '居民消费价格指数(上月=100)',
            'A01030G02': '食品烟酒类居民消费价格指数(上月=100)',
            'A01030G03': '衣着类居民消费价格指数(上月=100)',
            'A01030G04': '居住类居民消费价格指数(上月=100)',
            'A01030G05': '生活用品及服务类居民消费价格指数(上月=100)',
            'A01030G06': '交通通信类居民消费价格指数(上月=100)',
            'A01030G07': '教育文化娱乐类居民消费价格指数(上月=100)',
            'A01030G08': '医疗保健类居民消费价格指数(上月=100)',
            'A01030G09': '其他用品及服务类居民消费价格指数(上月=100)',
            'A01030G0A': '非食品居民消费价格指数(上月=100)',
            'A01030G0B': '消费品居民消费价格指数(上月=100)',
            'A01030G0C': '服务居民消费价格指数(上月=100)',
            'A01030G0D': '不包括食品和能源居民消费价格指数(上月=100)'
        },
        'A010301': {
            'A01030101': '居民消费价格指数(上月=100)',
            'A01030102': '食品烟酒类居民消费价格指数(上月=100)',
            'A01030103': '衣着类居民消费价格指数(上月=100)',
            'A01030104': '居住类居民消费价格指数(上月=100)',
            'A01030105': '生活用品及服务类居民消费价格指数(上月=100)',
            'A01030106': '交通和通信类居民消费价格指数(上月=100)',
            'A01030107': '教育文化和娱乐类居民消费价格指数(上月=100)',
            'A01030108': '医疗保健类居民消费价格指数(上月=100)',
            'A01030109': '其他用品和服务类居民消费价格指数(上月=100)'
        },
        'A010302': {
            'A01030201': '居民消费价格指数(上月=100)',
            'A01030202': '食品类居民消费价格指数(上月=100)',
            'A01030203': '烟酒及用品类居民消费价格指数(上月=100)',
            'A01030204': '衣着类居民消费价格指数(上月=100)',
            'A01030205': '家庭设备用品及维修服务类居民消费价格指数(上月=100)',
            'A01030206': '医疗保健和个人用品类居民消费价格指数(上月=100)',
            'A01030207': '交通和通讯类居民消费价格指数(上月=100)',
            'A01030208': '娱乐教育文化用品及服务类居民消费价格指数(上月=100)',
            'A01030209': '居住类居民消费价格指数(上月=100)'
        }
    }
    
    # 类常量：数据时期配置
    PERIOD_CONFIG = [
        {'zb': 'A01030G', 'sj': '2021-'},
        {'zb': 'A010301', 'sj': '2016-2020'},
        {'zb': 'A010302', 'sj': '-2015'}
    ]
    # 转换为字典便于查询
    PERIOD_MAP = {
        conf['sj']: conf['zb'] for conf in PERIOD_CONFIG
    }

    def __init__(self, db_config: dict = None):
        """
        初始化采集器
        :param db_config: 数据库配置字典，包含 host, user, password, database
        """
        self.session = requests.Session()
        self._init_session()
        self.db_config = db_config

    def _init_session(self):
        """配置请求会话"""
        retries = requests.adapters.Retry(
            total=3, backoff_factor=0.5, status_forcelist=[500, 502, 503, 504]
        )
        self.session.mount('https://', requests.adapters.HTTPAdapter(max_retries=retries))
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36 Edg/134.0.0.0'
        })

    @staticmethod
    def _generate_timestamp() -> int:
        """生成时间戳"""
        return int(time.time() * 1000)

    def get_available_periods(self) -> list:
        """获取可用的时间范围列表"""
        return list(self.PERIOD_MAP.keys())

    def _get_zbcode_by_sj(self, sj_range: str) -> str:
        """根据时间范围查找指标代码"""
        zb_code = self.PERIOD_MAP.get(sj_range)
        if not zb_code:
            available = self.get_available_periods()
            raise ValueError(
                f"无效的时间范围 '{sj_range}'。可用范围: {available}\n"
                "提示：2021- 表示2021年及以后，-2015 表示2015年及以前"
            )
        return zb_code
    
    def fetch_data(self, sj_range: str) -> Optional[pd.DataFrame]:
        """
        根据时间范围自动获取数据
        :param sj_range: 时间范围标识，如 '2021-', '2016-2020'
        """
        zb_code = self._get_zbcode_by_sj(sj_range)
        print(f"正在获取 {sj_range} 时期数据（指标代码：{zb_code}）...")
        
        try:
            # 构造请求参数
            params = {
                'm': 'QueryData',
                'dbcode': 'hgyd',
                'rowcode': 'zb',
                'colcode': 'sj',
                'wds': '[]',
                'dfwds': json.dumps([{"wdcode":"zb","valuecode":zb_code}, 
                                    {"wdcode":"sj","valuecode":sj_range}]),
                'k1': self._generate_timestamp()
            }
            
            # 发送请求
            response = self.session.get(
                'http://data.stats.gov.cn/easyquery.htm',
                params=params,
                timeout=10
            )
            response.raise_for_status()
            
            # 解析数据
            raw_data = response.json()
            print("成功获取并解析数据！")
            return self._parse_to_dataframe(raw_data, zb_code)
            
        except (requests.RequestException, json.JSONDecodeError, KeyError) as e:
            print(f"数据获取失败: {str(e)}")
            return None

    def _parse_to_dataframe(self, raw_data: dict, zb_code: str) -> pd.DataFrame:
        """解析JSON数据到DataFrame"""
        structured_data = []
        for node in raw_data['returndata']['datanodes']:
            # 添加防御性检查
            try:
                zb = next(wd['valuecode'] for wd in node['wds'] if wd['wdcode'] == 'zb')
                sj = next(wd['valuecode'] for wd in node['wds'] if wd['wdcode'] == 'sj')
                value = node['data']['strdata']
                structured_data.append({'zb': zb, 'sj': sj, 'value': value})
            except StopIteration:
                print(f"跳过无效节点：{node}")
                continue
        df = pd.DataFrame(structured_data).pivot(index='zb', columns='sj', values='value')
        
        # 时间列处理
        df.columns = pd.to_datetime(df.columns, format='%Y%m', errors='coerce')
        df = df.sort_index(axis=1, ascending=False)
        df.columns = df.columns.strftime('%Y-%m')
        
        # 添加中文名称列（保留原始代码）
        df = df.reset_index().rename(columns={'zb': 'zb_code'})
        df.insert(
            loc=1,  # 插入到zb列之后
            column='zb_name',
            value=df['zb_code'].map(self.CODE_MAPPING[zb_code])  # 映射中文名称
        )
        
        return df.set_index('zb_code')  # 恢复zb代码为索引
    
    @classmethod
    def save_to_excel(cls, df, filename):
        """
        将 DataFrame 保存到 Excel 文件
        """
        save_dir = "data"
        try:
            os.makedirs(save_dir, exist_ok=True)
            full_path = os.path.join(save_dir, filename)
            
            df.to_excel(full_path, index=False)
            
            print(f"Excel 文件已保存至：{os.path.abspath(full_path)}")
            return True
        except PermissionError:
            print(f"错误：文件 {filename} 被占用，请关闭后重试")
            return False
        except Exception as e:
            print(f"保存失败：{str(e)}")
            return False
